{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "xCF1aKktr7wM"
   },
   "source": [
    "To train this agent, click _Runtime_ and press _Run all_. Make sure you've enabled a free Tesla T4 GPU!\n",
    "\n",
    "<div class=\"align-center\">\n",
    "<a href=\"https://github.com/openpipe/art\"><img src=\"https://github.com/openpipe/art/raw/main/assets/ART_pill.png\" height=\"50\"></a>\n",
    "<a href=\"https://discord.gg/zbBHRUpwf4\"><img src=\"https://github.com/openpipe/art/raw/main/assets/Discord_pill.png\" height=\"50\"></a>\n",
    "<a href=\"https://art.openpipe.ai\"><img src=\"https://github.com/openpipe/art/raw/main/assets/Documentation_pill.png\" height=\"50\"></a>\n",
    "\n",
    "Questions? Join the Discord and ask away! For feature requests or to leave a star, visit our [Github](https://github.com/openpipe/art).\n",
    "\n",
    "</div>\n",
    "\n",
    "<a href=\"https://art.openpipe.ai/\"><img src=\"https://github.com/openpipe/art/raw/main/assets/Header_separator.png\" height=\"5\"></a>\n",
    "\n",
    "This notebook shows how to train a Qwen 2.5 3B model to play 2048. It will demonstrate how to set up a multi-turn agent, how to train it, and how to evaluate it.\n",
    "\n",
    "Completions and metrics will be logged to Weights & Biases.\n",
    "\n",
    "You will learn how to construct an [agentic environment](#Environment), how to define a [rollout](#Rollout), and how to run a [training loop](#Loop).\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "w4Y5hOMer7wO"
   },
   "source": [
    "### Installation\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "id": "KXsOex6lr7wO"
   },
   "outputs": [],
   "source": [
    "%%capture\n",
    "!uv pip install openpipe-art==0.3.11.post2 \"gql<4\" --prerelease allow --no-cache-dir"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "mqcaSOtFr7wN"
   },
   "source": [
    "### Environment Variables\n",
    "\n",
    "Later on in the notebook, we'll be creating a model that can automatically logs metrics to Weights & Biases and chat completions to Weave. In order to do so, you'll need to provide your Weights & Biases API key as an environment variable.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "id": "tyZkABN7r7wO"
   },
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "# Optional\n",
    "WANDB_API_KEY = \"\"\n",
    "if WANDB_API_KEY:\n",
    "    os.environ[\"WANDB_API_KEY\"] = WANDB_API_KEY"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "epnsrA1Ar7wO",
    "tags": [
     "environment"
    ]
   },
   "source": [
    "### Agentic Environment\n",
    "\n",
    "<a name=\"Environment\"></a>\n",
    "\n",
    "ART allows your agent to learn by interacting with its environment. In this example, we'll create an environment in which the agent can play 2048.\n",
    "\n",
    "Feel free to read as much or as little of this section's code as you'd like. The important thing to understand is that we're defining the rules of this agent's environment. In many cases, this will already be defined by the task you're trying to solve, but if you need to define a custom environment, this is how you do it.\n",
    "\n",
    "NOTE: To avoid OOM errors on a T4, we're reducing the winning value from 2048 to 128, which in turn reduces the minimum number of moves to win from 939 to 58.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "id": "SYWI9Ptlr7wO"
   },
   "outputs": [],
   "source": [
    "from dotenv import load_dotenv\n",
    "import random\n",
    "from typing import TypedDict\n",
    "from typing import Literal\n",
    "import string\n",
    "import xml.etree.ElementTree as ET\n",
    "\n",
    "load_dotenv()\n",
    "\n",
    "WINNING_VALUE = 128\n",
    "\n",
    "\n",
    "# Class that keeps track of state for a single game of 2048\n",
    "class TwentyFortyEightGame(TypedDict):\n",
    "    id: str\n",
    "    board: list[list[int | None]]\n",
    "\n",
    "\n",
    "# Randomly populates a cell on the board with a 2 or 4\n",
    "def populate_random_cell(game: TwentyFortyEightGame) -> None:\n",
    "    all_clear_coordinates = [\n",
    "        (i, j)\n",
    "        for i in range(len(game[\"board\"]))\n",
    "        for j in range(len(game[\"board\"][i]))\n",
    "        if game[\"board\"][i][j] is None\n",
    "    ]\n",
    "    random_clear_coordinates = random.choice(all_clear_coordinates)\n",
    "    # 90% chance to populate a 2, 10% chance to populate a 4\n",
    "    game[\"board\"][random_clear_coordinates[0]][random_clear_coordinates[1]] = (\n",
    "        2 if random.random() < 0.9 else 4\n",
    "    )\n",
    "\n",
    "\n",
    "# Generates a new game of 2048\n",
    "def generate_game(board_length: int = 4) -> TwentyFortyEightGame:\n",
    "    # random 6 character string\n",
    "    id = \"\".join(random.choices(string.ascii_letters + string.digits, k=6))\n",
    "    game = {\n",
    "        \"id\": id,\n",
    "        \"board\": [[None for _ in range(board_length)] for _ in range(board_length)],\n",
    "    }\n",
    "\n",
    "    # populate two random cells\n",
    "    populate_random_cell(game)\n",
    "    populate_random_cell(game)\n",
    "\n",
    "    return game\n",
    "\n",
    "\n",
    "# Renders the board in a human-readable format\n",
    "def render_board(game: TwentyFortyEightGame) -> str:\n",
    "    board = game[\"board\"]\n",
    "    # print something like this:\n",
    "    # _    | 2    | _    | 4\n",
    "    # 4    | 8    | 2    | 16\n",
    "    # 16   | 32   | 64   | 128\n",
    "    # _    | 2    | 2    | 4\n",
    "    # where _ is an empty cell\n",
    "\n",
    "    max_cell_width = max(\n",
    "        [len(str(cell)) for row in board for cell in row if cell is not None]\n",
    "    )\n",
    "\n",
    "    board_str = \"\"\n",
    "    for row in board:\n",
    "        # pad the cells with spaces to make them the same width\n",
    "        board_str += \"|\".join(\n",
    "            [\n",
    "                str(cell).rjust(max_cell_width)\n",
    "                if cell is not None\n",
    "                else \"_\".rjust(max_cell_width)\n",
    "                for cell in row\n",
    "            ]\n",
    "        )\n",
    "        board_str += \"\\n\"\n",
    "    return board_str\n",
    "\n",
    "\n",
    "# condense, privileging matches at the start of the sequence\n",
    "# sequences should be passed starting with cells that are the furthest in the direction in which the board is being condensed\n",
    "def condense_sequence(sequence: list[int | None]) -> list[int | None]:\n",
    "    condensed_sequence = []\n",
    "\n",
    "    gapless_sequence = [cell for cell in sequence if cell is not None]\n",
    "\n",
    "    i = 0\n",
    "    while i < len(gapless_sequence):\n",
    "        if (\n",
    "            i + 1 < len(gapless_sequence)\n",
    "            and gapless_sequence[i] == gapless_sequence[i + 1]\n",
    "        ):\n",
    "            condensed_sequence.append(gapless_sequence[i] * 2)\n",
    "            i += 2\n",
    "        else:\n",
    "            condensed_sequence.append(gapless_sequence[i])\n",
    "            i += 1\n",
    "\n",
    "    # pad the sequence with None at the end\n",
    "    return condensed_sequence + [None] * (4 - len(condensed_sequence))\n",
    "\n",
    "\n",
    "# Condenses the board in a given direction\n",
    "def condense_board(\n",
    "    game: TwentyFortyEightGame, direction: Literal[\"left\", \"right\", \"up\", \"down\"]\n",
    ") -> None:\n",
    "    if direction == \"left\":\n",
    "        for row in game[\"board\"]:\n",
    "            condensed_row = condense_sequence(row)\n",
    "            for i in range(len(row)):\n",
    "                row[i] = condensed_row[i]\n",
    "\n",
    "    if direction == \"right\":\n",
    "        for row in game[\"board\"]:\n",
    "            reversed_row = row[::-1]\n",
    "            # reverse the row before and after condensing\n",
    "            condensed_row = condense_sequence(reversed_row)[::-1]\n",
    "            for i in range(len(row)):\n",
    "                row[i] = condensed_row[i]\n",
    "\n",
    "    if direction == \"up\":\n",
    "        for col_index in range(len(game[\"board\"][0])):\n",
    "            column = [row[col_index] for row in game[\"board\"]]\n",
    "\n",
    "            condensed_column = condense_sequence(column)\n",
    "            for row_index in range(len(column)):\n",
    "                game[\"board\"][row_index][col_index] = condensed_column[row_index]\n",
    "\n",
    "    if direction == \"down\":\n",
    "        for col_index in range(len(game[\"board\"][0])):\n",
    "            column = [row[col_index] for row in game[\"board\"]]\n",
    "            reversed_column = column[::-1]\n",
    "            condensed_column = condense_sequence(reversed_column)[::-1]\n",
    "            for row_index in range(len(column)):\n",
    "                game[\"board\"][row_index][col_index] = condensed_column[row_index]\n",
    "\n",
    "\n",
    "# Applies an agent move to the game board\n",
    "def apply_agent_move(game: TwentyFortyEightGame, move_xml: str) -> None:\n",
    "    direction = None\n",
    "    # parse the move\n",
    "    try:\n",
    "        root = ET.fromstring(move_xml)\n",
    "        direction = root.text\n",
    "    except Exception:\n",
    "        raise ValueError(\"Invalid xml\")\n",
    "\n",
    "    if direction not in [\"left\", \"right\", \"up\", \"down\"]:\n",
    "        raise ValueError(\"Invalid direction\")\n",
    "\n",
    "    condense_board(game, direction)\n",
    "\n",
    "    populate_random_cell(game)\n",
    "\n",
    "\n",
    "# Returns the maximum cell value on the board\n",
    "def max_cell_value(game: TwentyFortyEightGame) -> int:\n",
    "    return max([cell for row in game[\"board\"] for cell in row if cell is not None])\n",
    "\n",
    "\n",
    "# Returns True if the game is finished\n",
    "def check_game_finished(game: TwentyFortyEightGame) -> bool:\n",
    "    if max_cell_value(game) >= WINNING_VALUE:\n",
    "        return True\n",
    "\n",
    "    # check if any cell is empty\n",
    "    if any(cell is None for row in game[\"board\"] for cell in row):\n",
    "        return False\n",
    "\n",
    "    return True\n",
    "\n",
    "\n",
    "# Returns the sum of all the cell values on the board\n",
    "def total_board_value(game: TwentyFortyEightGame) -> int:\n",
    "    return sum([cell for row in game[\"board\"] for cell in row if cell is not None])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "mwPZPzvbr7wP"
   },
   "source": [
    "### Creating a Model\n",
    "\n",
    "Now that we've defined the rules of our environment, we can create a model that will learn to play 2048. We'll use a Qwen 2.5 3B model for this example. The `name` parameter will be associated with a wandb run, and the `base_model` parameter is the model that we'll be training a LoRA on top of.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "tQhtIZFOr7wP"
   },
   "outputs": [],
   "source": [
    "import art\n",
    "from art.local import LocalBackend\n",
    "from dotenv import load_dotenv\n",
    "\n",
    "\n",
    "load_dotenv()\n",
    "\n",
    "random.seed(42)\n",
    "\n",
    "# Declare the model\n",
    "model = art.TrainableModel(\n",
    "    name=\"agent-002\",\n",
    "    project=\"2048-multi-turn\",\n",
    "    base_model=\"Qwen/Qwen2.5-3B-Instruct\",\n",
    ")\n",
    "# To run on a T4, we need to override some config defaults.\n",
    "model._internal_config = art.dev.InternalModelConfig(\n",
    "    init_args=art.dev.InitArgs(\n",
    "        max_seq_length=8192,\n",
    "    ),\n",
    "    engine_args=art.dev.EngineArgs(\n",
    "        enforce_eager=True,\n",
    "        gpu_memory_utilization=0.8,\n",
    "    ),\n",
    ")\n",
    "\n",
    "# Initialize the server\n",
    "backend = LocalBackend(\n",
    "    # Normally we don't want to run the server in-process, but for the output\n",
    "    # to show up properly on Google Colab we'll enable this.\n",
    "    in_process=True,\n",
    "    path=\"./.art\",\n",
    ")\n",
    "\n",
    "# Register the model with the local Backend (sets up logging, inference, and training)\n",
    "await model.register(backend)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "C8DY7vMir7wP",
    "tags": [
     "rollout"
    ]
   },
   "source": [
    "### Defining a Rollout\n",
    "\n",
    "<a name=\"Rollout\"></a>\n",
    "\n",
    "A rollout is a single episode of an agent performing its task. It generates one or more trajectories, which are lists of messages and choices.\n",
    "\n",
    "In this example, the rollout function generates a game of 2048, and the agent plays it until the game is finished. It then returns a trajectory which contains all the `system` and `user` messages presented to the agent, as well as all the `choices` that the agent made.\n",
    "\n",
    "When the game is finished the `reward` for the agent's performance is calculated based on the highest cell value on the board, which is then assigned to the trajectory.\n",
    "\n",
    "This rollout function will be called many times in parallel during each step of the training loop.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "id": "tXKhFp_Or7wP"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[36m\u001b[1mweave\u001b[0m: weave version 0.51.54 is available!  To upgrade, please run:\n",
      "\u001b[36m\u001b[1mweave\u001b[0m:  $ pip install weave --upgrade\n",
      "INFO:weave.trace.init_message:weave version 0.51.54 is available!  To upgrade, please run:\n",
      " $ pip install weave --upgrade\n",
      "\u001b[36m\u001b[1mweave\u001b[0m: Logged in as Weights & Biases user: arctic-fly.\n",
      "\u001b[36m\u001b[1mweave\u001b[0m: View Weave data at https://wandb.ai/art-tutorials/2048-multi-turn/weave\n",
      "INFO:weave.trace.init_message:Logged in as Weights & Biases user: arctic-fly.\n",
      "View Weave data at https://wandb.ai/art-tutorials/2048-multi-turn/weave\n"
     ]
    }
   ],
   "source": [
    "import art\n",
    "\n",
    "import math\n",
    "import requests\n",
    "from pydantic import BaseModel\n",
    "import weave\n",
    "\n",
    "if os.getenv(\"WANDB_API_KEY\", \"\"):\n",
    "    weave.init(model.project, settings={\"print_call_link\": False})\n",
    "\n",
    "\n",
    "class Scenario2048(BaseModel):\n",
    "    step: int\n",
    "\n",
    "\n",
    "@weave.op\n",
    "@art.retry(exceptions=(requests.ReadTimeout))\n",
    "async def rollout(model: art.Model, scenario: Scenario2048) -> art.Trajectory:\n",
    "    client = model.openai_client()\n",
    "    game = generate_game()\n",
    "\n",
    "    move_number = 0\n",
    "\n",
    "    trajectory = art.Trajectory(\n",
    "        messages_and_choices=[\n",
    "            {\n",
    "                \"role\": \"system\",\n",
    "                \"content\": \"You are an excellent 2048 player. Always choose the move most likely to lead to combine cells to eventually reach the number 2048. Optional moves are 'left', 'right', 'up', 'down'. Return your move as an XML object with a single property 'move', like so: <move>left</move>\",\n",
    "            }\n",
    "        ],\n",
    "        metadata={\n",
    "            \"game_id\": game[\"id\"],\n",
    "            \"notebook-id\": \"2048\",\n",
    "            \"step\": scenario.step,\n",
    "        },\n",
    "        reward=0,\n",
    "    )\n",
    "\n",
    "    while True:\n",
    "        trajectory.messages_and_choices.append(\n",
    "            {\"role\": \"user\", \"content\": render_board(game)}\n",
    "        )\n",
    "\n",
    "        try:\n",
    "            messages = trajectory.messages()\n",
    "            chat_completion = await client.chat.completions.create(\n",
    "                max_completion_tokens=128,\n",
    "                messages=messages,\n",
    "                model=model.name,\n",
    "                stream=False,\n",
    "            )\n",
    "        except Exception as e:\n",
    "            print(\"caught exception generating chat completion\", e)\n",
    "            raise e\n",
    "\n",
    "        choice = chat_completion.choices[0]\n",
    "        content = choice.message.content\n",
    "        assert isinstance(content, str)\n",
    "        trajectory.messages_and_choices.append(choice)\n",
    "\n",
    "        try:\n",
    "            apply_agent_move(game, content)\n",
    "            move_number += 1\n",
    "        except ValueError:\n",
    "            trajectory.reward = -1\n",
    "            break\n",
    "\n",
    "        if check_game_finished(game):\n",
    "            max_value = max_cell_value(game)\n",
    "            board_value = total_board_value(game)\n",
    "            trajectory.metrics[\"max_value\"] = max_value\n",
    "            trajectory.metrics[\"board_value\"] = board_value\n",
    "            trajectory.metrics[\"move_number\"] = move_number\n",
    "\n",
    "            # try to get as close to the winning value as possible\n",
    "            # otherwise, try to maximize number of high cells on board\n",
    "            # but above all else: WIN THE GAME!\n",
    "            if max_value < WINNING_VALUE:\n",
    "                # scale max value logarithmically between 0 for 2 and 1 for WINNING_VALUE\n",
    "                max_value_reward = (math.log(max_value, 2) - 1) / (\n",
    "                    math.log(WINNING_VALUE, 2) - 1\n",
    "                )\n",
    "                # scale board value logarithmically between 0 for 2 * 16 and 1 for WINNING_VALUE * 16\n",
    "                board_value_reward = (math.log(board_value, 2) - 1) / (\n",
    "                    math.log(WINNING_VALUE * 16, 2) - 1\n",
    "                )\n",
    "                # combine the two rewards, with max value having a higher weight\n",
    "                trajectory.reward = max_value_reward + (board_value_reward * 0.2)\n",
    "            else:\n",
    "                # double reward if the agent wins\n",
    "                trajectory.reward = 2\n",
    "            break\n",
    "\n",
    "    return trajectory"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "dr8DzLI4r7wP",
    "tags": [
     "loop"
    ]
   },
   "source": [
    "<a name=\"Loop\"></a>\n",
    "\n",
    "### Training Loop\n",
    "\n",
    "The training loop is where the magic happens. For each of the 10 steps defined below, the rollout function will be called 18 times in parallel. This means that 18 games will be played at once. Each game will produce a trajectory, which will be used to update the model.\n",
    "\n",
    "The `gather` step will wait for all of the trajectories to be generated, then it will delete all but the most recent checkpoint and train the model on the new trajectories.\n",
    "\n",
    "Inference will be blocked until the training is complete.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "id": "o-p5jaGur7wP"
   },
   "outputs": [],
   "source": [
    "for i in range(await model.get_step(), 10):\n",
    "    train_groups = await art.gather_trajectory_groups(\n",
    "        (\n",
    "            art.TrajectoryGroup(rollout(model, Scenario2048(step=i)) for _ in range(18))\n",
    "            for _ in range(1)\n",
    "        ),\n",
    "        pbar_desc=\"gather\",\n",
    "        max_exceptions=18,\n",
    "    )\n",
    "    await model.delete_checkpoints()\n",
    "    await model.train(\n",
    "        train_groups,\n",
    "        config=art.TrainConfig(learning_rate=1e-5),\n",
    "        # Lowering the logprob_calculation_chunk_size is a memory saving measure\n",
    "        # to allow longer sequences (up to 8192 tokens) to be processed on a T4.\n",
    "        _config={\"logprob_calculation_chunk_size\": 8},\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "EmSndgphr7wP"
   },
   "source": [
    "### Using the Model\n",
    "\n",
    "Just like that, you've trained an agent to play 2048! Now it's time to use your model outside of ART, in the wild! The easiest way to do that is to load it from disk, where it was saved after each training step, and either run inference on it locally or upload it to a central hub like HuggingFace.\n",
    "\n",
    "Check out the code below for small demo of the model you just trained playing 2048!\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "id": "T5mMSMJ9r7wP",
    "outputId": "febd4699-83d3-4c07-e67e-5acdaf34d345"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loading model from .art/2048-multi-turn/models/agent-002/0010\n",
      "\n",
      "==((====))==  Unsloth 2025.5.1: Fast Qwen2 patching. Transformers: 4.51.3. vLLM: 0.8.5.post1.\n",
      "   \\\\   /|    NVIDIA H100 PCIe. Num GPUs = 1. Max memory: 79.179 GB. Platform: Linux.\n",
      "O^O/ \\_/ \\    Torch: 2.6.0+cu124. CUDA: 9.0. CUDA Toolkit: 12.4. Triton: 3.2.0\n",
      "\\        /    Bfloat16 = TRUE. FA [Xformers = 0.0.29.post2. FA2 = False]\n",
      " \"-____-\"     Free license: http://github.com/unslothai/unsloth\n",
      "Unsloth: Fast downloading is enabled - ignore downloading bars which are red colored!\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "The attention mask is not set and cannot be inferred from input because pad token is same as eos token. As a consequence, you may observe unexpected behavior. Please pass your input's `attention_mask` to obtain reliable results.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "move 10\n",
      "board:\n",
      "_|2|4|2\n",
      "_|_|4|4\n",
      "2|_|4|2\n",
      "_|_|_|_\n",
      "\n",
      "agent move: <move>down</move>\n",
      "updated board:\n",
      "_|_|_|_\n",
      "4|_|_|2\n",
      "_|_|4|4\n",
      "2|2|8|2\n",
      "\n",
      "\n",
      "move 20\n",
      "board:\n",
      " _| _| 2| 4\n",
      " _| _| _|16\n",
      " _| _| 2| 8\n",
      " _| 4| 2| 8\n",
      "\n",
      "agent move: <move>up</move>\n",
      "updated board:\n",
      " _| 4| 4| 4\n",
      " _| _| 2|16\n",
      " _| 2| _|16\n",
      " _| _| _| _\n",
      "\n",
      "\n",
      "move 30\n",
      "board:\n",
      " 2| _| _| _\n",
      " 2| 4| _| _\n",
      " 2| 8| 2| _\n",
      " 2| 8| 4|32\n",
      "\n",
      "agent move: <move>down</move>\n",
      "updated board:\n",
      " _| _| _| 2\n",
      " _| _| _| _\n",
      " 4| 4| 2| _\n",
      " 4|16| 4|32\n",
      "\n",
      "\n",
      "move 40\n",
      "board:\n",
      " _| _|16| 4\n",
      " _| 8|32| 8\n",
      " _| _|16| 4\n",
      " _| 2| _| _\n",
      "\n",
      "agent move: <move>down</move>\n",
      "updated board:\n",
      " 2| _| _| _\n",
      " _| _|16| 4\n",
      " _| 8|32| 8\n",
      " _| 2|16| 4\n",
      "\n",
      "\n",
      "move 50\n",
      "board:\n",
      " _| _| _| 4\n",
      " 2| _|16| 4\n",
      " _| 4| 8|32\n",
      " 2|16| 8|16\n",
      "\n",
      "agent move: <move>up</move>\n",
      "updated board:\n",
      " 4| 4|16| 8\n",
      " _|16|16|32\n",
      " _| _| _|16\n",
      " _| _| _| 2\n",
      "\n",
      "\n",
      "move 60\n",
      "board:\n",
      " 8| _| _| _\n",
      " 2| 4| 2| _\n",
      "32|64| _| 2\n",
      " 4| 8| 8| _\n",
      "\n",
      "agent move: <move>down</move>\n",
      "updated board:\n",
      " 8| _| _| _\n",
      " 2| 4| 2| _\n",
      "32|64| 2| _\n",
      " 4| 8| 8| 2\n",
      "\n",
      "\n",
      "move 70\n",
      "board:\n",
      " _|32| 4| 4\n",
      " 4|16| 4|64\n",
      " _| 2|16| 2\n",
      " 2| _| 2| 4\n",
      "\n",
      "agent move: <move>down</move>\n",
      "updated board:\n",
      " 2| _| _| 4\n",
      " _|32| 8|64\n",
      " 4|16|16| 2\n",
      " 2| 2| 2| 4\n",
      "\n",
      "game finished in 78 moves\n",
      "game won! 💪\n",
      "final board:\n",
      "\n",
      "  _|  _|  _|  _\n",
      "  2|  8|  _|  _\n",
      " 16|128|  2|  2\n",
      "  4|  4|  4|  8\n",
      "\n",
      "max value: 128\n",
      "board value: 178\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "from unsloth import FastLanguageModel\n",
    "\n",
    "\n",
    "# example: .art/2048-multi-turn/models/001/0003\n",
    "lora_model_path = (\n",
    "    f\".art/{model.project}/models/{model.name}/{await model.get_step():04d}\"\n",
    ")\n",
    "\n",
    "print(f\"loading model from {lora_model_path}\\n\")\n",
    "\n",
    "peft_model, tokenizer = FastLanguageModel.from_pretrained(\n",
    "    model_name=lora_model_path,\n",
    "    max_seq_length=16384,\n",
    "    dtype=torch.bfloat16,\n",
    "    load_in_4bit=True,\n",
    ")\n",
    "FastLanguageModel.for_inference(peft_model)\n",
    "\n",
    "game = generate_game()\n",
    "move_number = 0\n",
    "\n",
    "messages = [\n",
    "    {\n",
    "        \"role\": \"system\",\n",
    "        \"content\": \"You are an excellent 2048 player. Always choose the move most likely to lead to combine cells to eventually reach the number 2048. Optional moves are 'left', 'right', 'up', 'down'. Return your move as an XML object with a single property 'move', like so: <move>left</move>\",\n",
    "    },\n",
    "]\n",
    "\n",
    "while not check_game_finished(game):\n",
    "    rendered_board = render_board(game)\n",
    "    messages.append({\"role\": \"user\", \"content\": rendered_board})\n",
    "\n",
    "    inputs = tokenizer.apply_chat_template(\n",
    "        messages, return_tensors=\"pt\", add_generation_prompt=True\n",
    "    ).to(\"cuda\")\n",
    "\n",
    "    content = \"\"\n",
    "\n",
    "    def get_completion() -> str:\n",
    "        with torch.no_grad():\n",
    "            outputs = peft_model.generate(\n",
    "                input_ids=inputs,\n",
    "                max_new_tokens=100,\n",
    "                do_sample=True,\n",
    "                temperature=0.7,\n",
    "                top_p=0.9,\n",
    "            )\n",
    "            return tokenizer.decode(\n",
    "                outputs[0][inputs.shape[1] :], skip_special_tokens=True\n",
    "            )\n",
    "\n",
    "    try:\n",
    "        content = get_completion()\n",
    "    except Exception as e:\n",
    "        print(\"caught exception generating chat completion\", e)\n",
    "        raise e\n",
    "\n",
    "    messages.append({\"role\": \"assistant\", \"content\": content})\n",
    "\n",
    "    try:\n",
    "        apply_agent_move(game, content)\n",
    "        move_number += 1\n",
    "    except ValueError:\n",
    "        raise ValueError(f\"Invalid move on move {move_number}: {content}\")\n",
    "\n",
    "    # print the board every 10 moves\n",
    "    if move_number % 10 == 0:\n",
    "        print(f\"\\nmove {move_number}\")\n",
    "        print(f\"board:\\n{rendered_board}\")\n",
    "        print(f\"agent move: {content}\")\n",
    "        print(f\"updated board:\\n{render_board(game)}\")\n",
    "\n",
    "\n",
    "print(f\"game finished in {move_number} moves\")\n",
    "\n",
    "max_value = max_cell_value(game)\n",
    "board_value = total_board_value(game)\n",
    "\n",
    "if max_value >= WINNING_VALUE:\n",
    "    print(\"game won! 💪\")\n",
    "else:\n",
    "    print(\"game lost! 😢\")\n",
    "\n",
    "\n",
    "print(f\"final board:\\n\\n{render_board(game)}\")\n",
    "print(f\"max value: {max_value}\")\n",
    "print(f\"board value: {board_value}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "4m3eqX_Vr7wQ"
   },
   "source": [
    "<div class=\"align-center\">\n",
    "<a href=\"https://github.com/openpipe/art\"><img src=\"https://github.com/openpipe/art/raw/notebooks/assets/ART_pill.png\" height=\"50\"></a>\n",
    "<a href=\"https://discord.gg/zbBHRUpwf4\"><img src=\"https://github.com/openpipe/art/raw/notebooks/assets/Discord_pill.png\" height=\"50\"></a>\n",
    "<a href=\"https://openpipe.ai/blog/art-e-mail-agent\"><img src=\"https://github.com/openpipe/art/raw/main/assets/ART_E_pill.png\" height=\"50\"></a>\n",
    "\n",
    "Questions? Join the Discord and ask away! For feature requests or to leave a star, visit our [Github](https://github.com/openpipe/art).\n",
    "\n",
    "</div>\n"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "provenance": []
  },
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
